import akshare as ak
import pandas as pd
import os
import time
from datetime import datetime, timedelta
from typing import List

# 配置请求头
ak.tool.user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36"

def get_all_stock_list() -> List[str]:
    """获取全市场股票代码（沪深A股+北交所）"""
    # 沪市A股
    sh_stocks = ak.stock_info_sh_name_code(symbol="主板A股")["证券代码"].tolist()
    # 深市A股
    sz_stocks = ak.stock_info_sz_name_code(symbol="A股列表")["A股代码"].tolist()
    # 北交所股票（新版接口）
    try:
        bj_df = ak.stock_bj_spot_em()
        bj_stocks = bj_df["symbol"].tolist()  # 确认列名为symbol
    except Exception as e:
        print(f"北交所数据获取失败: {str(e)}")
        bj_stocks = []
    
    return [
        *[f"sh{code}" for code in sh_stocks if len(code) == 6],
        *[f"sz{code}" for code in sz_stocks if len(code) == 6],
        *bj_stocks
    ]

def get_history_data(code: str) -> pd.DataFrame:
    """获取个股历史数据"""
    try:
        # 沪深A股处理
        if code.startswith(("sh", "sz")):
            pure_code = code[2:]
            df = ak.stock_zh_a_hist(
                symbol=pure_code,
                period="daily",
                start_date=(datetime.now() - timedelta(days=365)).strftime("%Y%m%d"),
                end_date=datetime.now().strftime("%Y%m%d"),
                adjust="qfq"
            )
        # 北交所处理
        elif code.startswith("bj"):
            df = ak.stock_bj_daily(symbol=code)
        else:
            return pd.DataFrame()
            
        return df if not df.empty else df
    except Exception as e:
        print(f"获取 {code} 失败: {str(e)}")
        return pd.DataFrame()

if __name__ == "__main__":
    os.makedirs('stock_data', exist_ok=True)
    codes = get_all_stock_list()
    print(f"共获取 {len(codes)} 只股票")
    
    for i, code in enumerate(codes, 1):
        df = get_history_data(code)
        if not df.empty:
            df.to_csv(f"stock_data/{code}.csv", index=False)
            print(f"进度 {i}/{len(codes)} 已保存 {code}")
        else:
            print(f"进度 {i}/{len(codes)} 无数据 {code}")
        # time.sleep(0.01)

    print("全部完成！")